package com.raven.play.datalink.jdbc;

import akka.Done;
import akka.stream.Attributes;
import akka.stream.Inlet;
import akka.stream.Shape;
import akka.stream.SinkShape;
import akka.stream.stage.AbstractInHandler;
import akka.stream.stage.GraphStageLogic;
import akka.stream.stage.GraphStageWithMaterializedValue;
import akka.stream.stage.TimerGraphStageLogic;
import lombok.extern.slf4j.Slf4j;
import scala.Tuple2;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.time.Duration;
import java.util.ArrayDeque;
import java.util.Queue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;

/**
 * JDBC 输出源
 */
@Slf4j
public abstract class JdbcSinkGraphStage<I> extends GraphStageWithMaterializedValue<SinkShape<I>, CompletionStage<Done>> {

    private final Inlet<I> inlet = Inlet.create("DatabaseSink.in");
    private final SinkShape<I> shape = new SinkShape<>(inlet);
    private final DataSource dataSource;
    private final String insertSql;
    private final StatementBuilder<I> statementBuilder;
    private BatchExecutionOptions batchExecOpts;


    public JdbcSinkGraphStage(DataSource dataSource, String insertSql, StatementBuilder<I> statementBuilder, BatchExecutionOptions batchExecOpts) {
        this.dataSource = dataSource;
        this.insertSql = insertSql;
        this.statementBuilder = statementBuilder;
        this.batchExecOpts = batchExecOpts;
    }

    @Override
    public SinkShape<I> shape() {
        return shape;
    }

    @Override
    public Tuple2<GraphStageLogic, CompletionStage<Done>> createLogicAndMaterializedValue(Attributes inheritedAttributes) throws Exception {
        CompletableFuture<Done> promise = new CompletableFuture<>();
        return new Tuple2<>(new JdbcSinkGraphStageLogic(shape, promise), promise);
    }

    class JdbcSinkGraphStageLogic extends TimerGraphStageLogic {

        // 用于批量保存的时间控制器
        private static final String TIMER_KEY = "BatchTimer";
        // 缓冲区
        private final Queue<I> buffer;
        private Connection connection;
        private PreparedStatement preparedStatement;
        private final CompletableFuture<Done> promise;

        public JdbcSinkGraphStageLogic(Shape shape, CompletableFuture<Done> promise) {
            super(shape);
            this.buffer = new ArrayDeque<>(batchExecOpts.getBatchSize());
            this.promise = promise;

            this.setHandler(inlet, new AbstractInHandler() {
                @Override
                public void onPush() throws Exception {
                    I input = grab(inlet);
                    buffer.add(input);
                    // 如果到达了批次的临界值, 则执行保存
                    if (buffer.size() >= batchExecOpts.getBatchSize()) {
                        executeBatch();
                    }
                    pull(inlet);
                }

            });
        }

        @Override
        public void preStart() throws Exception {
            if (batchExecOpts.needTimer()) {
                schedulePeriodically(TIMER_KEY, Duration.ofMillis(batchExecOpts.getIntervalMs()));
            }

            try {
                connection = dataSource.getConnection();
                preparedStatement = connection.prepareStatement(insertSql);
            } catch (SQLException e) {
                log.error(e.getMessage(), e);
                failStage(e);
            }

            // 开始流处理前拉取第一个元素
            pull(inlet);
        }

        @Override
        public void postStop() throws Exception {
            try {
                executeBatch();
                if (preparedStatement != null) preparedStatement.close();
                if (connection != null) connection.close();
            } catch (SQLException e) {
                // 处理关闭资源时的异常
                log.info(e.getMessage(), e);
            }

            promise.complete(Done.getInstance());
        }

        @Override
        public void onTimer(Object timerKey) throws Exception {
            // 如果达到了间隔市场, 则执行保存
            if (timerKey.equals(TIMER_KEY)) {
                try {
                    executeBatch();
                } catch (SQLException e) {
                    log.error(e.getMessage(), e);
                    failStage(e);
                }
            }
        }

        private void executeBatch() throws Exception {
            if (buffer.isEmpty()) {
                return;
            }
            while (!buffer.isEmpty()) {
                I input = buffer.poll();
                statementBuilder.assembly(preparedStatement, input);
                preparedStatement.addBatch();
            }
            preparedStatement.executeBatch();
            if (!connection.getAutoCommit()) {
                connection.commit();
            }
        }
    }
}
